import os
import os.path as osp
import argparse
import tensorflow as tf
import torch
import torchvision
import torchvision.transforms as transforms
import numpy as np
from utils import get_test_dataset


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', type=str, default=osp.join('data', 'mnist_net.pb'))
    parser.add_argument('--dataset', type=str, choices=['mnist', 'cifar10'])
    args = parser.parse_args()

    with tf.gfile.FastGFile(args.model, 'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())

    with tf.Graph().as_default() as graph:
        tf.import_graph_def(graph_def, name='')

    input_tensor = graph.get_tensor_by_name('input:0')
    output_tensor = graph.get_tensor_by_name('output:0')

    _, test_loader = get_test_dataset(args.dataset)

    for x, label in testloader:
        image = np.ascontiguousarray(x.numpy())
        with tf.Session(graph=graph) as sess:
            output_vals = sess.run(output_tensor, feed_dict={input_tensor: image})
        print('label={}, y={}'.format(label, int(np.argmax(np.array(output_vals).squeeze(), axis=0))))
